#!/usr/bin/env python  
#-*- coding:utf-8 _*-  
""" 
@author:hello_life 
@license: Apache Licence 
@file: albert.py 
@time: 2022/05/05
@software: PyCharm 
description:
"""
import torch
import torch.nn as nn

from transformers import AlbertModel

class Albert_NER(nn.Module):
    def __init__(self,config):
        super(Albert_NER,self).__init__()
        self.config=config
        self.albert=AlbertModel.from_pretrained(config.pretrained_model_path)
        self.classifier=nn.Linear(384,config.num_class)

    def forward(self,batch):
        output=self.albert(**batch)
        output=self.classifier(output.last_hidden_state)
        return output

    def forward_with_loss(self,token_ids,masks,y):
        output = self.albert(token_ids, masks)
        output = self.classifier(output.last_hidden_state)
        loss_fct=nn.CrossEntropyLoss()

        loss=loss_fct(output.view(-1,self.config.num_class),y.view(-1))

        return output,loss